// Copyright (c) 2020 Presto Labs Pte. Ltd.
// Author: donggu

#include "coin2/exchange/base/feed/ws/ws_subscriber.h"

#include <memory>

#include <boost/asio/spawn.hpp>

#include "coin2/base/log.h"

namespace coin2::exchange::base::feed {

void WsSubscriber::SubscribeAsync() {
  for (auto& it : connection_configs_) {
    AddForSymbolGroupAndWorker(&it);

    if (it.mea_config->unauthorized_limit.max_subscription) {
      CHECK_LE_THROW(it.symbols.size(), it.mea_config->unauthorized_limit.max_subscription.value())
          << "[WsSubscriber] exceeds unauthorized max_subscription for " <<
          it.mea_config->mea.String();
    }
  }

  sub_running_ = true;
  std::weak_ptr<WsSubscriber> self(shared_from_this());
  boost::asio::spawn(strand_, [self](boost::asio::yield_context yield) {
    if (auto s = self.lock()) {
      s->Subscribe(std::move(yield));
    } else {
      LOG(ERROR) << "boost::asio handler object (WsSubscriber) does not exist";
    }
  });

  if (spec_.periodic_force_reconnect_minute > 0) {
    scheduler_->AddPeriodicTask(
        [self] {
          if (auto s = self.lock()) {
            s->ForceReconnect();
          } else {
            LOG(ERROR) << "boost::asio handler object (WsSubscriber) does not exist";
          }
        },
        spec_.periodic_force_reconnect_minute * 60 * 1'000'000'000LL,
        "force reconnect");
  }

  if (config_.reconnect_slowest() || config_.monitor_connection()) {
      std::weak_ptr<WsSubscriber> self(shared_from_this());
      scheduler_->AddPeriodicTask(
        [self] {
          if (auto s = self.lock()) {
            LOG(INFO) << "Reconnect Monitor checking...";
            const auto ret = s->latency_manager_.RaceConnectionLatency();
            for (const auto& config : ret) {
              s->latency_manager_.RemoveLatencyExtractor(config);
              if (!config->symbols.empty()) {
                s->ForceReconnectWorker(config);
              } else {
                std::string info = config->topic + "_" + \
                    config->url.name + "_" + \
                    std::to_string(config->connection_id);
                LOG(INFO) << info << "has no symbols to subscribe";
              }
            }
          } else {
            LOG(ERROR) << "WsConnection handler is already destroyed";
          }
        },
        5 * 60 * 1'000'000'000LL,  // 5 min
        "reconnect_monitor",
        5 * 60 * 1'000'000'000LL);  // 5 min
    }
}

bool WsSubscriber::BanProduct(const std::vector<const IProduct*>& products) {
  std::map<WsConnectionConfig*, std::vector<std::string>> connection_configs;
  FindAndRemoveSymbolFromConfig(products, &connection_configs);
  for (auto& [connection_config, removed_symbols] : connection_configs) {
    std::weak_ptr<WsSubscriber> self(shared_from_this());
    if (connection_config->symbols.size() > 0) {
      boost::asio::post(
          strand_, 
          [self, connection_config, removed_symbols] {
            if (auto s = self.lock()) {
              std::string info = connection_config->topic + "_" + \
                  connection_config->url.name + "_" + \
                  std::to_string(connection_config->connection_id);
              LOG(INFO) << "Ban " << removed_symbols.size() << " products from " << info;
              s->UnsubscribeForSymbolGroupAndWorker(connection_config, removed_symbols);
            }
          });
    } else {
      boost::asio::post(
          strand_, 
          [self, connection_config] {
            if (auto s = self.lock()) {
              std::string info = connection_config->topic + "_" + \
                  connection_config->url.name + "_" + \
                  std::to_string(connection_config->connection_id);
              LOG(INFO) << "Ban all products from " << info;
              s->RemoveForSymbolGroupAndWorker(connection_config);
            }
          });
    }
  }
  return false; 
}

void WsSubscriber::FindAndRemoveSymbolFromConfig(
    const std::vector<const IProduct*> products,
    std::map<WsConnectionConfig*, std::vector<std::string>>* connection_configs) {
  std::set<std::string> native_symbols;
  for (const auto& p : products) {
    native_symbols.insert(p->feed_native_symbol());
  }

  for (auto& config : connection_configs_) {
    std::vector<std::string> removed_config_symbols;
    auto& symbols = config.symbols;
    for (auto iter = symbols.begin(); iter != symbols.end(); ) {
      if (native_symbols.find(*iter) != native_symbols.end()) {
        removed_config_symbols.push_back(*iter);
        iter = symbols.erase(iter);
      } else {
        ++iter;
      }
    }

    if (!removed_config_symbols.empty()) {
      connection_configs->insert({&config, removed_config_symbols});
    }
  }
}

}  // namespace coin2::exchange::base::feed
